#include <arpa/inet.h>

#include "dp_trie_ip.h"

namespace dp {
dp_trie_ip_node::dp_trie_ip_node()
    : data_(NULL)
{
    children_[0] = nullptr;
    children_[1] = nullptr;
}

dp_trie_ip_node::~dp_trie_ip_node() { }

dp_trie_ip::dp_trie_ip() { root_ = std::make_shared<dp_trie_ip_node>(); }

dp_trie_ip::~dp_trie_ip() { }

bool dp_trie_ip::insert(const std::string& cidr, void* value)
{
    std::vector<unsigned int> ip;
    int mask = 0;
    if (!parse_cidr(cidr, ip, mask, NULL)) {
        return false;
    }
    if (mask <= 0 || NULL == value) {
        return false;
    }

    auto node = root_;
    int bit_index = 0; // 位索引
    for (int i = 0; i < mask; ++i) {
        int part_index = i / 32; // 每32位处理一部分
        int bit = (ip[part_index] >> (31 - bit_index)) & 1;
        if (node->children_[bit] == nullptr) {
            node->children_[bit] = std::make_shared<dp_trie_ip_node>();
        }
        node = node->children_[bit];
        bit_index = (bit_index + 1) % 32;
    }
    node->data_ = value;
    return true;
}

void* dp_trie_ip::search(const std::string& ip_str)
{
    std::vector<unsigned int> ip;
    if (ip_str.empty() || false == ip2integer(ip_str, ip, NULL)) {
        return NULL;
    }
    auto node = root_;
    void* ret = NULL;
    int bit_index = 0;
    int max = ip.size() * 32;
    for (int i = 0; i < max; ++i) { // IPv6是128位
        int part_index = i / 32;
        int bit = (ip[part_index] >> (31 - bit_index)) & 1;
        if (node->children_[bit] != nullptr) {
            node = node->children_[bit];
            if (node->data_) {
                ret = node->data_;
            }
        } else {
            break;
        }
        bit_index = (bit_index + 1) % 32;
    }
    return ret;
}

bool dp_trie_ip::ip2integer(const std::string& ip_str, std::vector<unsigned int>& ip, bool* bv6)
{
    bool bret = false;
    if (ip_str.find(':') != std::string::npos) { // 检测是否为IPv6
        if (bv6) {
            *bv6 = true;
        }
        std::vector<unsigned int> result(4, 0); // IPv6分为4个32位部分
        struct sockaddr_in6 in6;
        if (1 == inet_pton(AF_INET6, ip_str.c_str(), &in6.sin6_addr)) {
            const unsigned char* addr_bytes = in6.sin6_addr.s6_addr;
            for (int i = 0; i < 4; ++i) {
                // 逐个将 4 个字节组合成一个 32 位无符号整数
                result[i] = (addr_bytes[i * 4] << 24) | (addr_bytes[i * 4 + 1] << 16) | (addr_bytes[i * 4 + 2] << 8) | (addr_bytes[i * 4 + 3]);
            }
            ip = result;
            bret = true;
        }
    } else { // IPv4
        std::vector<unsigned int> result(1, 0);
        struct sockaddr_in in;
        if (1 == inet_pton(AF_INET, ip_str.c_str(), &in.sin_addr)) {
            const unsigned char* addr_bytes = reinterpret_cast<const unsigned char*>(&in.sin_addr);
            // 将4字节组合成一个 32 位无符号整数，存储到 result[0]
            result[0] = (addr_bytes[0] << 24) | (addr_bytes[1] << 16) | (addr_bytes[2] << 8) | (addr_bytes[3]);
            ip = result;
            bret = true;
        }
    }
    return bret;
}

bool dp_trie_ip::parse_cidr(const std::string& cidr, std::vector<unsigned int>& ip,
    int& mask, bool* bv6)
{
    size_t pos = cidr.find('/');
    std::string ipStr = cidr.substr(0, pos);
    mask = std::stoi(cidr.substr(pos + 1));
    return ip2integer(ipStr, ip, bv6);
}

} // dp
